#!/usr/bin/env python
# coding=utf-8
# @FileName :excel_util.py
# @Time     :2022/2/10 13:16
# @Author   :ZXC
# @Desc     :excel读写工具类
from pathlib import Path

import xlrd
import openpyxl
from loguru import logger
from openpyxl.styles import PatternFill
from xlrd import XLRDError
from xlrd.sheet import Sheet


class ExcelReadUtil(object):
    """
    Excel读取数据工具类
    """

    def __init__(self, file_path=None, file_contents=None, sheet_name=None):
        """
        Excel读取数据工具类初始化方法
        Args:
            file_path: 文件路径
            file_contents: 文件内容
            sheet_name: sheet页名称,不填写默认读取第一个sheet页
        """
        self.file_path = file_path
        self.file_contents = file_contents
        self.sheet_name = sheet_name
        self.workbook = None
        self.sheet = None
        logger.debug(f'正在初始化Excel读取数据工具类，文件路径【{file_path}】,文件内容{file_contents},sheet页名称{sheet_name}')
        try:
            self.workbook = xlrd.open_workbook(filename=self.file_path, file_contents=file_contents)
        except FileNotFoundError:
            raise FileNotFoundError(f'Excel文件【{file_path}、未找到')
        except TypeError:
            raise TypeError(f'Excel文件【{file_path}、类型错误')
        except Exception as e:
            raise Exception(f'Excel文件解析失败，原因：{e}')

        if self.sheet_name:
            logger.debug(f'Excel文件初始化完成，加载{sheet_name}页签')
            try:
                self.sheet = self.workbook.sheet_by_name(sheet_name=self.sheet_name)
                logger.debug(f'Excel文件初始化完成，加载{sheet_name}页签')
            except XLRDError as e:
                raise XLRDError(f'Excel文件Sheet页{self.sheet_name}读取失败，原因：{e}')
        else:
            logger.debug(f'Excel文件初始化完成，加载第一个页签')
            self.sheet = self.workbook.sheet_by_index(sheetx=0)

    def get_sheets(self):
        """
        获取excel的所有sheet
        Returns: xlrd.sheet.Sheet object列表

        """
        logger.debug('获取excel的所有sheet')
        value = self.workbook.sheets()
        logger.debug(f'获取excel的所有sheet完成：{value}')
        return value

    def get_sheet_names(self):
        """
        获取excel所有sheet的名称
        Returns:

        """
        logger.debug(f'获取excel所有sheet的名称')
        value = self.workbook.sheet_names()
        logger.debug(f'获取excel所有sheet的名称完成：{value}')
        return value

    def get_sheet_name(self, sheet: Sheet = None):
        """
        获取sheet名称
        Args:
            sheet:

        Returns:

        """
        if sheet:
            ws = sheet
        else:
            ws = self.sheet
        logger.debug(f'获取{ws.name}页的名称')
        value = ws.name
        logger.debug(f'获取{ws.name}页的名称完成：{value}')
        return value

    def get_nrows(self, sheet: Sheet = None):
        """
        获取Sheet页总行数
        Args:
            sheet:

        Returns:

        """
        if sheet:
            ws = sheet
        else:
            ws = self.sheet
        logger.debug(f'获取{ws.name}页的总行数')
        value = ws.nrows
        logger.debug(f'获取{ws.name}页的总行数完成：{value}')
        return value

    def get_ncols(self):
        """
        获取Sheet页总列数
        Returns:

        """
        logger.debug(f'获取{self.sheet.name}页的总列数')
        value = self.sheet.ncols
        logger.debug(f'获取{self.sheet.name}页的总列数完成：{value}')
        return value

    def get_title(self, nrow=0):
        """
        读取Excel表头
        Args:
            nrow: 表头行，默认第一行为表头信息

        Returns: 数据列表

        """
        logger.debug(f'读取Excel表头(第{nrow}行)')
        try:
            data = [value for value in self.sheet.row_values(nrow)]
            logger.debug(f'读取Excel表头第{nrow}行完成：{data}')
            return data
        except Exception as e:
            raise Exception(f'读取Excel数据失败，原因：{e}')

    def get_datas(self, start_row=None, end_row=None, start_col=None, end_col=None):
        """
        读取Excel数据，存放在一个二维列表中返回
        Args:
            start_row: 起始行
            end_row: 结束行
            start_col: 起始列
            end_col: 线束列

        Returns: 数据列表

        """
        start_row = start_row or 0
        end_row = end_row or self.get_nrows()
        start_col = start_col or 0
        end_col = end_col or self.get_ncols()

        logger.debug(f'读取Excel数据,第{start_row}到{end_row}行,第{start_col}到{end_col}列')

        try:
            datas = list()
            for nrow in range(start_row, end_row):
                data = [value for value in self.sheet.row_values(nrow)]
                datas.append(data[start_col:end_col])
            logger.debug(f'读取读取Excel数据,第{start_row}到{end_row}行,第{start_col}到{end_col}列数据完成：{datas}')
            return datas
        except Exception as e:
            raise Exception(f'读取Excel数据失败，原因：{e}')

    def get_dict_fmt_datas(self, start_row=None, end_row=None, start_col=None, end_col=None):
        """
        获取Sheet页数据，并将数据转换成dict，第一行数据为key
        Args:
            start_row: 起始行
            end_row: 结束行
            start_col: 起始列
            end_col: 线束列

        Returns:

        """
        logger.debug(f'获取Sheet页数据，并将数据转换成dict,第{start_row}到{end_row}行,第{start_col}到{end_col}列')
        keys = self.get_title()
        values = self.get_datas(start_row, end_row, start_col, end_col)
        res = [dict(zip(keys, value)) for value in values]
        logger.debug(f'获取Sheet页数据，并将数据转换成dict,第{start_row}到{end_row}行,第{start_col}到{end_col}列完成：{res}')
        return res

    def get_dict_fmt_datas_generator(self, start_row=None, end_row=None, start_col=None, end_col=None):
        """
        获取Sheet页数据，并将数据转换成dict，第一行数据为key
        此方式为列表生成器，效率比较高，少占用内存
        Args:
            start_row: 起始行
            end_row: 结束行
            start_col: 起始列
            end_col: 线束列

        Returns: 生成器对象

        """
        start_row = start_row or 1
        end_row = end_row or self.get_nrows()
        start_col = start_col or 0
        end_col = end_col or self.get_ncols()

        logger.debug(f'获取Sheet页数据，并将数据转换成dict对象的列表生成器,第{start_row}到{end_row}行,第{start_col}到{end_col}列')
        keys = self.get_title()
        res = (dict(zip(keys, self.sheet.row_values(nrow))) for nrow in range(start_row, end_row))
        logger.debug(f'获取Sheet页数据，列表生成器对象为{res}')
        return res

    def get_datas_ignore_ncols(self, ncols):
        """
        读取Excel数据，忽略指定列
        Args:
            ncols: 需要忽略的列，类型为列表

        Returns: 数据列表

        """
        logger.debug(f'读取Excel数据，忽略指定列{ncols}')
        try:
            datas = list()
            for nrow in range(self.get_nrows()):
                data = list()
                for ncol in range(self.get_ncols()):
                    if ncol not in ncols:
                        data.append(self.sheet.cell_value(rowx=nrow, colx=ncol))
                datas.append(data)
            logger.debug(f'读取Excel数据，忽略指定列{ncols}完成{datas}')
            return datas
        except Exception as e:
            raise Exception(f'读取Excel数据失败，原因：{e}')

    def get_datas_ignore_nrows(self, nrows):
        """
        读取Excel数据，忽略指定行
        Args:
            nrows: 需要忽略的行，类型为列表

        Returns: 数据列表

        """
        logger.debug(f'读取Excel数据，忽略指定行{nrows}')
        try:
            datas = list()
            for nrow in range(self.get_nrows()):
                if nrow not in nrows:
                    data = [value for value in self.sheet.row_values(nrow)]
                    datas.append(data)
            logger.debug(f'读取Excel数据，忽略指定行{nrows}完成{datas}')
            return datas
        except Exception as e:
            raise Exception(f'读取Excel数据失败，原因：{e}')

    def get_data_by_nrow(self, nrow):
        """
        读取Excel指定行数据
        Args:
            nrow: 需要读取数据的行

        Returns: 数据列表

        """
        logger.debug(f'读取Excel指定行{nrow}的数据')
        try:
            data = [value for value in self.sheet.row_values(nrow)]
            logger.debug(f'读取Excel指定行{nrow}的数据完成：{data}')
            return data
        except Exception as e:
            raise Exception(f'读取Excel数据失败，原因：{e}')

    def get_data_by_ncol(self, ncol):
        """
        读取Excel指定列数据
        Args:
            ncol: 需要读取数据的列

        Returns: 数据列表

        """
        logger.debug(f'读取Excel指定列{ncol}的数据')
        try:
            data = [value for value in self.sheet.col_values(ncol)]
            logger.debug(f'读取Excel指定列{ncol}的数据完成；{data}')
            return data
        except Exception as e:
            raise Exception(f'读取Excel数据失败，原因：{e}')

    def get_data_by_cell(self, row: int, col: int):
        """
        读取Excel指定单元格数据
        Args:
            row: 需要读取数据的单元格的行，首行为0
            col: 需要读取数据的单元格的列，首列为0

        Returns:

        """
        logger.debug(f'读取Excel指定单元格({row},{col})的数据')
        try:
            data = self.sheet.cell_value(row - 1, col - 1)
            logger.debug(f'读取Excel指定单元格({row},{col})的数据完成：{data}')
            return data
        except Exception as e:
            raise Exception(f'读取Excel数据失败，原因：{e}')


class ExcelWriteUtil(object):
    """
    Excel写入数据工具类
    """

    def __init__(self, file_path=None, sheet_name=None):
        """
        Excel写入数据工具类初始化方法
        Args:
            file_path: 文件路径
            sheet_name: sheet页名称
        """
        self.file_path = file_path
        self.sheet_name = sheet_name
        self.workbook = None
        self.sheet = None

        if not file_path:
            file_path = ''

        if Path(file_path).exists() and Path(file_path).is_file():
            logger.debug(f'正在初始化Excel写入数据工具类,文件路径【{file_path}、,Sheet页名称{sheet_name}')
            self.workbook = openpyxl.load_workbook(filename=self.file_path)
            if self.sheet_name:
                try:
                    self.sheet = self.workbook[self.sheet_name]
                except KeyError as e:
                    raise KeyError(f'Excel文件Sheet页加载失败，原因：{e}')
        else:
            logger.debug(f'正在初始化Excel写入数据工具类,新建一个文件')
            self.workbook = openpyxl.Workbook()

        if not self.sheet:
            self.sheet = self.workbook.active
        logger.debug(f'初始化Excel写入工具类完成，文件{self.workbook},Sheet页{self.sheet}')

    def create_sheet(self, title, index=0):
        """
        创建sheet页
        Args:
            title:
            index:

        Returns:

        """
        logger.debug(f'正在创建Sheet页，名称{title}，位置{index}')
        self.sheet = self.workbook.create_sheet(title, index)
        logger.debug(f'创建Sheet页完成，名称{title}，位置{index}')

    def remove_sheet_by_name(self, name):
        """
        根据sheet名称删除
        Args:
            name:

        Returns:

        """
        logger.debug(f'正在移除指定Sheet页，名称{name}')
        self.workbook.remove(self.workbook[name])
        logger.debug(f'移除指定Sheet页完成，名称{name}')

    def remove_all_sheet(self):
        """
        删除所有sheet页
        Returns:

        """
        logger.debug(f'正在移除所有Sheet页')
        for name in self.workbook.sheetnames:
            self.workbook.remove(self.workbook[name])
        logger.debug(f'移除所有sheet页完成')

    def append_data_by_row(self, datas):
        """
        按行追加内容
        Args:
            datas:

        Returns:

        """
        logger.debug(f'正在添加数据：{datas}')
        self.sheet.append(iterable=datas)
        logger.debug(f'添加数据完成：{datas}')

    def update_data_by_cell(self, row, column, value):
        """
        修改单元格数据
        Args:


        Returns:

        """
        logger.debug(f'正在更新单元格({row},{column})数据：{value}')
        self.sheet.cell(row, column, value)
        logger.debug(f'更新单元格({row},{column})数据完成：{value}')

    def write_data_by_cell(self, row, column, value):
        """
        单元格写入数据
        Args:
            row: 行
            column: 列
            value: 数据

        Returns:

        """
        self.update_data_by_cell(row, column, value)

    def fill_cell_color(self, row, col, color):
        """
        单元格填充背景色
        Args:
            row: 行
            col: 列
            color: 背景色

        Returns:

        """
        logger.debug(f'正在给单元格({row},{col})填充背景色：{color}')
        pattern_fill = PatternFill("solid", color)
        self.sheet.cell(row, col).fill = pattern_fill
        logger.debug(f'给单元格({row},{col})填充背景色完成：{color}')

    def save(self, filename=None):
        """
        保存文件
        Args:
            filename: 文件名称

        Returns:

        """
        logger.debug(f'正在保存文件{filename}')
        filename = filename if filename else self.file_path
        self.workbook.save(filename=filename)
        logger.debug(f'保存文件完成{filename}')


if __name__ == '__main__':
    from config.settings import DATA_DIR

    filepath = DATA_DIR.joinpath('test_case.xlsx')
    # print(filepath)

    wb = ExcelReadUtil(file_path=filepath)
    # print(wb.get_sheets())
    # print(wb.get_sheet_names())
    # print(wb.get_sheet_name())
    # print(wb.get_nrows())
    # print(wb.get_ncols())
    # print(wb.get_title())
    # print(wb.get_datas())
    # print(wb.get_dict_fmt_datas(1))
    # print(wb.get_dict_fmt_datas_generator(1).__next__())
    for row_data in wb.get_dict_fmt_datas_generator():
        print(row_data)
    # print(wb.get_datas_ignore_ncols([0,1]))
    # print(wb.get_datas_ignore_nrows([0,1,2,3,4]))
    # print(wb.get_data_by_nrow(1))
    # print(wb.get_data_by_ncol(1))
    # print(wb.get_data_by_cell(3, 2))
    # nwb = ExcelWriteUtil(file_path=filepath)
    # nwb.create_sheet('zxc')
    # nwb.remove_sheet_by_name('zxc')
    # nwb.remove_all_sheet()
    # nwb.create_sheet('zxc')
    # nwb.append_data_by_row({'A':'a','C':'c'})
    # nwb.update_data_by_cell(1,1,'zxc')

    # nwb.save()
